package find;

import java.util.HashMap;
import java.util.List;
import java.util.Stack;

/**
 * @author Liaorun
 */
public class UnionFind {


    public static class Element<V> {
        public V value;


        public Element(V value) {
            this.value = value;
        }


        @Override
        public int hashCode() {
            return super.hashCode();
        }

        @Override
        public boolean equals(Object obj) {
            return super.equals(obj);
        }
    }

    public static class UnionFindSet<V> {

        public HashMap<V, Element<V>> elementHashMap;

        /**
         * key:某个元素，V:该元素的母节点
         */
        public HashMap<Element<V>, Element<V>> fatherMap;

        /**
         * K：每个树的顶节点 V: 树的节点数
         */
        public HashMap<Element<V>, Integer> sizeMap;

        public UnionFindSet(List<V> list) {
            elementHashMap = new HashMap<>();
            fatherMap = new HashMap<>();
            sizeMap = new HashMap<>();

            for (V value : list) {
                Element<V> element = new Element<>(value);
                elementHashMap.put(value, element);
                fatherMap.put(element, element);
                sizeMap.put(element, 1);
            }
        }

        /**
         * 给点一个节点，向上一直找，把树顶的元素返回
         *
         * @param element 一个节点元素
         * @return 树顶的元素 or null
         */
        private Element<V> findHead(Element<V> element) {

            Stack<Element<V>> path = new Stack<>();

            // 循环一直找到该节点所在树的顶点
            while (element != fatherMap.get(element)) {
                path.push(element);
                element = fatherMap.get(element);
            }

            // 将树变成变成扁平的，方便后面查询
            while (!path.isEmpty()) {
                fatherMap.put(path.pop(), element);
            }

            return element;
        }


        /**
         * 两个节点是否是同一个集合
         *
         * @param a 一个元素
         * @param b 一个元素
         * @return true:是
         */
        public boolean isSameSet(V a, V b) {
            if (elementHashMap.containsKey(a) && elementHashMap.containsKey(b)) {
                // 所在树的顶是一样的就是同一个集合
                return findHead(elementHashMap.get(a)) == findHead(elementHashMap.get(b));
            }

            return false;
        }

        /**
         * 合并两个节点所在的两个树
         *
         * @param a 一个节点元素
         * @param b 一个节点元素
         */
        public void union(V a, V b) {

            if (elementHashMap.containsKey(a) && elementHashMap.containsKey(b)) {
                // a, b都是添加过的节点

                // a元素所在树的顶
                Element<V> aF = findHead(elementHashMap.get(a));
                // b元素所在树的顶
                Element<V> bF = findHead(elementHashMap.get(b));


                // 判断树顶是否相同
                if (aF != bF) {
                    // 不同合并两棵树

                    // 大树的顶点
                    Element<V> big = sizeMap.get(aF) >= sizeMap.get(bF) ? aF : bF;
                    // 小树的顶点
                    Element<V> small = big == aF ? bF : aF;

                    // 小树变成大树的子树，小树顶指向大树的顶
                    fatherMap.put(small, big);

                    // 修改记录的树的大小
                    sizeMap.put(big, sizeMap.get(aF) + sizeMap.get(bF));
                    sizeMap.remove(small);
                }
            }
        }
    }
}
